/**
* @file Basic_Opera.c
* 
* @author W. P. Yao, ywp@pku.edu.cn
* 
* @date 2015-12-10
*
* @brief 
* 1. void Convolution        --- get phi and psi
* 2. void Divergence         --- get rho_x, rho_y; omega_x, omega_y
* 3. void Value_u_v          --- get u and v from omega
* 4. void Transpose_Matrix   --- transpose martrix during MPI
* 5. void FFT_R2C_2D         --- Fourier from R to C, Add_NUM long
* 6. void FFT_C2R_2D         --- Fourier from C to R, Add_NUM long
* 7. void FFT_R2C_For_Input  --- Fourier from R to C, Num     long
* 8. void FFT_C2R_For_Output --- Fourier from C to R, Num     long
*/
#include "Basic_Opera.h"

/**
* @Name: Convolution 
*
* @Param: A
* @Param: B
*
* @Brief: 1. cal u*rho_x   + v*rho_y,   Fourier to C, get phi
* 		  2. cal u*omega_x + v*omega_y, Fourier to C, get psi
*/
void Convolution(fftw_complex *A, fftw_complex *B)
{
	Divergence(C_rho, C_rho_x, C_rho_y);				///cal div
	Divergence(C_omega, C_omega_x, C_omega_y);

	FFT_C2R_2D(C_u, R_Add_u);							///Fourier C2R
	FFT_C2R_2D(C_v, R_Add_v);
	FFT_C2R_2D(C_rho_x, R_Add_rho_x);
	FFT_C2R_2D(C_rho_y, R_Add_rho_y);
	FFT_C2R_2D(C_omega_x, R_Add_omega_x);
	FFT_C2R_2D(C_omega_y, R_Add_omega_y);

	for(k = 0; k < Add_M*Add_NUM; k++){                 ///cal phi and psi in R
		R_Add_phi[k][0] = R_Add_u[k][0]*R_Add_rho_x[k][0]   +
						  R_Add_v[k][0]*R_Add_rho_y[k][0];
		R_Add_phi[k][1] = 0;

		R_Add_psi[k][0] = R_Add_u[k][0]*R_Add_omega_x[k][0] +
						  R_Add_v[k][0]*R_Add_omega_y[k][0];
		R_Add_psi[k][1] = 0;
	}

	FFT_R2C_2D(R_Add_phi, A);							///Fourier R2C
	FFT_R2C_2D(R_Add_psi, B);

}

/**
* @Name: Divergence 
*
* @Param: IN
* @Param: OUT_X
* @Param: OUT_Y
*/
void Divergence(fftw_complex *IN, fftw_complex *OUT_X, fftw_complex *OUT_Y)
{
	double Restrict;
	double tmp_k;
	double k1, k2, tmp;

	for(j = 0; j < M; j++){
		tmp = j + M*myrank;
		if(tmp < NUM/2){
			k1 = tmp;
		}
		else{
			k1 = tmp - NUM;
		}
		if(k1 == -NUM/2){
			k1 = 0;
		}
		tmp_k = k1;
		
		for(i = 0; i < NUM; i++){		
			if(i < NUM/2){
				k2 = i;
			}
			else{
				k2 = i - NUM;
			}
			if(k2 == -NUM/2){
				k2 = 0;
			}
			tmp_k = k2;

			OUT_X[index][0] = -IN[index][1]*k1;
			OUT_X[index][1] =  IN[index][0]*k1;
			OUT_Y[index][0] = -IN[index][1]*k2;
			OUT_Y[index][1] =  IN[index][0]*k2;
		}
	}
}



/**
* @Name: Value_u_v 
*
* @Param: IN
* @Param: OUT_U
* @Param: OUT_V
* @Brief: cal u and v from omega
*/
void Value_u_v(fftw_complex *IN, fftw_complex *OUT_U, fftw_complex *OUT_V)
{
	double Restrict;
	double tmp_k;
	double tmp, TMP;
	double k1, k2;
	for(j = 0; j < M; j++){
		tmp = j + M*myrank;
		if(tmp < NUM/2){
			k1 = tmp;
		}
		else{
			k1 = tmp - NUM;
		}
		if(k1 == -NUM/2){
			k1 = 0;
		}
		tmp_k = k1;
		
		for(i = 0; i < NUM; i++){		
			if(i < NUM/2){
				k2 = i;
			}
			else{
				k2 = i - NUM;
			}
			if(k2 == -NUM/2){
				k2 = 0;
			}
			TMP = k1*k1 + k2*k2;

			if(TMP == 0){
				OUT_U[index][0] = 0;
				OUT_U[index][1] = 0;
				OUT_V[index][0] = 0;
				OUT_V[index][1] = 0;
			}
			else{
				OUT_U[index][0] =  k2*IN[index][1]/TMP;
				OUT_U[index][1] = -k2*IN[index][0]/TMP;
				OUT_V[index][0] = -k1*IN[index][1]/TMP;
				OUT_V[index][1] =  k1*IN[index][0]/TMP;
			}
		}
	}
}


/**
* @Name: Transpose_Matrix 
*
* @Param: A
* @Param: B
* @Param: N_x
* @Param: M_y
* @Param: N_y
* @Param: M_x
*/
void Transpose_Matrix(fftw_complex *A, fftw_complex *B, int N_x, int M_y, int N_y, int M_x)
{
	MPI_Type_vector(N_x, N_y, N_y, MPI_DOUBLE, &matrix_type);   ///create datatype, a matrix of M*N  ???
	MPI_Type_commit(&matrix_type);								///commit
	MPI_Type_vector(N_y, 1, N_x, MPI_DOUBLE, &column_type);		
	MPI_Type_commit(&column_type);
	MPI_Type_hvector(N_x, 1, sizeof(double), column_type, &transposed_matrix_type);
	MPI_Type_commit(&transposed_matrix_type);

	double *send_buffer, *recv_buffer;
	send_buffer = (double *)calloc(N_x*M_y, sizeof(double));
	recv_buffer = (double *)calloc(N_y*M_x, sizeof(double));

	for(k = 0; k < sum_rank; k++){
		for(j = 0; j < N_x; j++){
			for(i = 0; i < N_y; i++){
				send_buffer[k*N_x*N_y + j*N_y + i] = A[j*M_y + i + k*N_y][0];	/// send real part
			}
		}
	}
	MPI_Barrier(MPI_COMM_WORLD);
	MPI_Alltoall(send_buffer,1,matrix_type,recv_buffer,1,transposed_matrix_type,MPI_COMM_WORLD);
	MPI_Barrier(MPI_COMM_WORLD);
	for(k = 0; k < sum_rank; k++){
		for(j = 0; j < N_y; j++){
			for(i = 0; i < N_x; i++){
				B[j*M_x + i + k*N_x][0] = recv_buffer[k*N_y*N_x + j*N_x + i];	/// recv real part
			}
		}
	}
	
	for(k = 0; k < sum_rank; k++){
		for(j = 0; j < N_x; j++){
			for(i = 0; i < N_y; i++){
				send_buffer[k*N_x*N_y + j*N_y + i] = A[j*M_y + i + k*N_y][1];	/// send imaginary part
			}
		}
	}
	MPI_Barrier(MPI_COMM_WORLD);
	MPI_Alltoall(send_buffer,1,matrix_type,recv_buffer,1,transposed_matrix_type,MPI_COMM_WORLD);
	MPI_Barrier(MPI_COMM_WORLD);
	for(k = 0; k < sum_rank; k++){
		for(j = 0; j < N_y; j++){
			for(i = 0; i < N_x; i++){
				B[j*M_x + i + k*N_x][1] = recv_buffer[k*N_y*N_x + j*N_x + i];	/// recv imaginary part
			}
		}
	}
	
	free(send_buffer);
	free(recv_buffer);
	MPI_Type_free(&matrix_type);
	MPI_Type_free(&column_type);
	MPI_Type_free(&transposed_matrix_type);

}


/**
* @Name: FFT_R2C_2D 
*
* @Param: A, Add_M*Add_NUM long
* @Param: B, M*NUM long
*
* @Brief: Fourier from R to C
*/
void FFT_R2C_2D(fftw_complex *A, fftw_complex *B)
{
	int I_tmp1 = NUM/2;
	fftw_complex * temp1, *temp2;
	temp1 = (fftw_complex*)fftw_malloc((Add_M*NUM)*sizeof(fftw_complex));
	temp2 = (fftw_complex*)fftw_malloc((M*Add_NUM)*sizeof(fftw_complex));

	for(j = 0; j < Add_M; j++){                              /// Fourier in x, R to C
		for(i = 0; i < Add_NUM; i++){
			in[i][0] = A[C_INDEX[j][i]][0];
			in[i][1] = 0;
		}
		fftw_execute(fftw_fp);
		for(i = 0; i < I_tmp1; i++){
			temp1[index][0] = out[i][0]/Add_NUM;
			temp1[index][1] = out[i][1]/Add_NUM;
		}
		for(i = I_tmp1; i < NUM; i++){
			temp1[index][0] = out[i+Inset_N][0]/Add_NUM;
			temp1[index][1] = out[i+Inset_N][1]/Add_NUM;
		}
	}

	Transpose_Matrix(temp1, temp2, Add_M, NUM, M, Add_NUM);

	for(j = 0; j < M; j++){                              /// Fourier in y, R to C
		for(i = 0; i < Add_NUM; i++){
			in[i][0] = temp2[C_INDEX[j][i]][0];
			in[i][1] = temp2[C_INDEX[j][i]][1];
		}
		fftw_execute(fftw_fp);
		for(i = 0; i < I_tmp1; i++){
			B[index][0] = out[i][0]/Add_NUM;
			B[index][1] = out[i][1]/Add_NUM;
		}
		for(i = I_tmp1; i < NUM; i++){
			B[index][0] = out[i+Inset_N][0]/Add_NUM;
			B[index][1] = out[i+Inset_N][1]/Add_NUM;
		}
	}

	fftw_free(temp1);
	fftw_free(temp2);
}


/**
* @Name: FFT_C2R_2D 
*
* @Param: A, M*NUM
* @Param: B, Add_M*Add_NUM
*
* @Brief: Fourier from C to R
*/
void FFT_C2R_2D(fftw_complex *A, fftw_complex *B)
{
	int I_tmp1, I_tmp2;
	fftw_complex * temp1, *temp2;
	temp2 = (fftw_complex*)fftw_malloc((Add_M*NUM)*sizeof(fftw_complex));
	temp1 = (fftw_complex*)fftw_malloc((M*Add_NUM)*sizeof(fftw_complex));

	I_tmp1 = NUM/2;
	I_tmp2 = I_tmp1 + Inset_N;
	
	for(j = 0; j < M; j++){                              /// Inverse Fourier in y, C to R
		for(i = 0; i < I_tmp1; i++){
			in[i][0] = A[index][0];
			in[i][1] = A[index][1];
		}
		for(i = I_tmp1; i < NUM; i++){
			in[i + Inset_N][0] = A[index][0];
			in[i + Inset_N][1] = A[index][1];
		}
		for(i = I_tmp1; i < I_tmp2; i++){
			in[i][0] = 0;
			in[i][1] = 0;
		}
		fftw_execute(fftw_bp);
		for(i = 0; i < Add_NUM; i++){
			temp1[C_INDEX[j][i]][0] = out[i][0];
			temp1[C_INDEX[j][i]][1] = out[i][1];
		}
	}

	Transpose_Matrix(temp1, temp2, Add_M, NUM, M, Add_NUM);

	for(j = 0; j < M; j++){                              /// Fourier in y, R to C
		for(i = 0; i < Add_NUM; i++){
			in[i][0] = temp2[C_INDEX[j][i]][0];
			in[i][1] = temp2[C_INDEX[j][i]][1];
		}
		fftw_execute(fftw_fp);
		for(i = 0; i < I_tmp1; i++){
			B[index][0] = out[i][0]/Add_NUM;
			B[index][1] = out[i][1]/Add_NUM;
		}
		for(i = I_tmp1; i < NUM; i++){
			B[index][0] = out[i+Inset_N][0]/Add_NUM;
			B[index][1] = out[i+Inset_N][1]/Add_NUM;
		}
	}

	fftw_free(temp1);
	fftw_free(temp2);
}


/**
* @Name: FFT_R2C_For_Input 
*
* @Param: A, M*NUM
* @Param: B, M*NUM
*/
void FFT_R2C_For_Input(fftw_complex *A, fftw_complex *B)
{
	fftw_complex * temp1, *temp2;
	temp1 = (fftw_complex*)fftw_malloc((M*NUM)*sizeof(fftw_complex));
	temp2 = (fftw_complex*)fftw_malloc((M*NUM)*sizeof(fftw_complex));

	for(j = 0; j < M; j++){                              /// Fourier in x, R to C
		for(i = 0; i < NUM; i++){
			in[i][0] = A[index][0];
			in[i][1] = A[index][1];
		}
		fftw_execute(fftw_fip);
		for(i = 0; i < NUM; i++){
			temp1[index][0] = Out[i][0]/NUM;
			temp1[index][1] = Out[i][1]/NUM;
		}
	}

	Transpose_Matrix(temp1, temp2, M, NUM, M, NUM);

	for(j = 0; j < M; j++){                              /// Fourier in y, R to C
		for(i = 0; i < NUM; i++){
			in[i][0] = temp2[index][0];
			in[i][1] = temp2[index][1];
		}
		fftw_execute(fftw_fip);
		for(i = 0; i < NUM; i++){
			B[index][0] = Out[i][0]/NUM;
			B[index][1] = Out[i][1]/NUM;
		}
	}

	fftw_free(temp1);
	fftw_free(temp2);
}


/**
* @Name: FFT_C2R_For_Output 
*
* @Param: A, M*NUM
* @Param: B, M*NUM
*/
void FFT_C2R_For_Output(fftw_complex *A, fftw_complex *B)
{
	fftw_complex * temp1, *temp2;
	temp1 = (fftw_complex*)fftw_malloc((M*NUM)*sizeof(fftw_complex));
	temp2 = (fftw_complex*)fftw_malloc((M*NUM)*sizeof(fftw_complex));

	for(j = 0; j < M; j++){                              /// Fourier in x, R to C
		for(i = 0; i < NUM; i++){
			in[i][0] = A[index][0];
			in[i][1] = A[index][1];
		}
		fftw_execute(fftw_bop);
		for(i = 0; i < NUM; i++){
			temp1[index][0] = Out[i][0];
			temp1[index][1] = Out[i][1];
		}
	}

	Transpose_Matrix(temp1, temp2, M, NUM, M, NUM);

	for(j = 0; j < M; j++){                              /// Fourier in y, R to C
		for(i = 0; i < NUM; i++){
			in[i][0] = temp2[index][0];
			in[i][1] = temp2[index][1];
		}
		fftw_execute(fftw_bop);
		for(i = 0; i < NUM; i++){
			B[index][0] = Out[i][0];
			B[index][1] = Out[i][1];
		}
	}

	fftw_free(temp1);
	fftw_free(temp2);
}
